#pragma once

#include "KVDatabase.h"
#include <boost/archive/text_oarchive.hpp>
#include <boost/archive/text_iarchive.hpp>
#include <boost/serialization/unordered_map.hpp>
#include <shared_mutex>
#include <unordered_map>
#include <string>
#include <stdexcept>
#include <mutex>

template <typename Key, typename Value>
class ThreadSafeKVDB : public KVDatabase<Key, Value> {
private:
    std::unordered_map<Key, Value> data_;
    mutable std::shared_mutex mutex_;

public:
    ThreadSafeKVDB() = default;

    // 插入或更新键值对
    bool put(const Key& key, const Value& value) override {
        std::unique_lock lock(mutex_);
        data_[key] = value;
        return true;
    }

    // 查询键值对
    bool get(const Key& key, Value& value) const override {
        std::shared_lock lock(mutex_);
        auto it = data_.find(key);
        if (it != data_.end()) {
            value = it->second;
            return true;
        }
        return false;
    }

    // 判断键是否存在
    bool exists(const Key& key) const override {
        std::shared_lock lock(mutex_);
        return data_.count(key) == 1;
    }

    // 删除键值对
    bool remove(const Key& key) override {
        std::unique_lock lock(mutex_);
        return data_.erase(key) > 0;
    }

    // 序列化：将当前数据库状态写入 data_str
    bool serialize(std::string& data_str) const override {
        std::unique_lock lock(mutex_);
        try {
            std::ostringstream oss;
            boost::archive::text_oarchive oa(oss);
            oa << data_;
            data_str = oss.str();
            return true;
        }
        catch (...) {
            return false;
        }
    }

    // 反序列化：从 data_str 加载数据库状态
    bool deserialize(const std::string& data_str) override {
        std::unique_lock lock(mutex_);
        try {
            std::istringstream iss(data_str);
            boost::archive::text_iarchive ia(iss);
            ia >> data_;
            return true;
        }
        catch (...) {
            return false;
        }
    }

    // 获取当前键值对数量
    size_t size() const override {
        std::shared_lock lock(mutex_);
        return data_.size();
    }
};